from tools.recoder_args import DefaultArgs

from bdtime import Time
from bdtime import tt
import os
from tools import my_cv2_utils
import pandas as pd
import numpy as np
from tools import ocr_utils
import re
from tqdm import tqdm
import cv2
import json

from bdtime import show_ls
# from tools.screen_recorder import run as run__screen_recorder, ScreenRecorder
# from tools.utils import get_public_attrs
# import threading


def is_fraction(s):
    pattern = r'^\d+/\d+$'
    match = re.match(pattern, s)
    return bool(match)


class Coordinates:
    is_complete_coordinates = [(210, 335), (260, 350)]
    extinct_enemy_coordinates = [(65, 104), (92, 115)]


class RewardAnalysis:
    def __init__(self, operation_name, img_dir_name,
                 record_dir_path=DefaultArgs.record_path.recorded_state_dir_path, i_range: tuple = None):
        self.record_dir_path = record_dir_path
        self.screen_dir_path = os.path.join(record_dir_path, 'screen_image')

        img_dir_path = os.path.join(self.screen_dir_path, operation_name, img_dir_name)
        assert os.path.exists(img_dir_path), f'img_dir_path[{img_dir_path}]不存在?'
        self.img_dir_path = img_dir_path

        key_file_path = os.path.join(record_dir_path, 'key', operation_name)
        assert os.path.exists(key_file_path)
        self.key_file_path = key_file_path

        mouse_file_path = os.path.join(record_dir_path, 'rel_mouse', operation_name)
        assert os.path.exists(mouse_file_path)
        self.mouse_file_path = mouse_file_path

        with open(key_file_path, 'r') as f:
            key_ls = json.load(f)
        key_ls = [i for i in key_ls if 'i' in i]

        with open(mouse_file_path, 'r') as f:
            mouse_ls = json.load(f)
        mouse_ls = [i for i in mouse_ls if 'i' in i]

        self.analysis_result_save_dir = os.path.join(
            DefaultArgs.record_path.dir_path__analysis_result, operation_name, img_dir_name)
        os.makedirs(self.analysis_result_save_dir, exist_ok=True)
        self.output_path__reward_state_info_df = os.path.join(self.analysis_result_save_dir, 'reward_state_info_df.csv')

        self._reward_state_info = []
        self.reward_state_info_df = pd.DataFrame(self._reward_state_info, columns=['game_state', 'extinct_enemy', 'remain_enemy'])

        _img_ls = os.listdir(img_dir_path)
        _img_ls = [i for i in _img_ls if i.endswith('jpg') or i.endswith('png')]
        get_create_time = lambda filename: os.path.getctime(os.path.join(img_dir_path, filename))
        ordered_img_ls = sorted(_img_ls, key=get_create_time)

        if isinstance(i_range, tuple):
            assert len(i_range) == 2, 'i_range的长度必须为2!'
            ordered_img_ls = ordered_img_ls[i_range[0]: i_range[1]]
            key_ls = key_ls[i_range[0]: i_range[1]]
            mouse_ls = mouse_ls[i_range[0]: i_range[1]]

            self.output_path__reward_state_info_df = os.path.join(self.analysis_result_save_dir, f'reward_state_info_df__{i_range[0]}_to_{i_range[1]}.csv')

        img_path_ls = [os.path.join(img_dir_path, i) for i in ordered_img_ls]
        info_ls = [os.path.splitext(i)[0].split('__')[1:] for i in ordered_img_ls]

        dtype = {'i': np.int32, 'time': np.float32}
        df = pd.DataFrame(info_ls, columns=['i', 'time']).astype(dtype)
        col = pd.Series(ordered_img_ls)
        df['file_name'] = col
        col = pd.Series(img_path_ls)
        df['file_path'] = col

        col = pd.Series(mouse_ls)
        df['mouse'] = col
        col = pd.Series(key_ls)
        df['key'] = col

        # show_ls(key_ls[:5])
        # show_ls(mouse_ls[:5])

        # len(key_ls), len(mouse_ls), len(img_path_ls)

        self.key_ls = key_ls
        self.mouse_ls = mouse_ls

        self.df = df
        self.img_path_ls = img_path_ls
        self.ordered_img_ls = ordered_img_ls
        self.info_ls = info_ls

        self.total_frames = len(img_path_ls)
        self.key_total_frames = len(key_ls)

        self.coordinates = Coordinates()

        self.tt = Time()

        self.cache_info_dc = {}

        self.total_enemy = -1
        self.get_total_enemy()

    def get_mid_val(self, game_state, extinct_enemy):
        if game_state == 0:
            mid_val = extinct_enemy
        elif game_state == -1:
            mid_val = -1
        elif game_state == 1:
            mid_val = self.total_enemy
        else:
            raise ValueError('game_state取值错误! 必须在[-1, 0, 1]之中!')
        return mid_val

    def get_img_i(self, i):
        img_i = my_cv2_utils.read_image(self.img_path_ls[i])
        return img_i

    def show_img(self, img_i):
        if isinstance(img_i, int):
            img_i = self.get_img_i(img_i)
        my_cv2_utils.show_img(img_i)

    def get_info(self, img_i, show=False):
        _img_i = None
        if isinstance(img_i, int):
            _img_i = img_i
            if img_i in self.cache_info_dc:
                game_state, extinct_enemy, total_enemy = self.cache_info_dc.get(img_i)
                return game_state, extinct_enemy, total_enemy

            img_i = self.get_img_i(img_i)

        extinct_enemy, total_enemy = self._get_extinct_enemy_and_total_enemy(img_i)

        is_complete = self._get_is_complete(img_i)
        is_start = not is_complete and extinct_enemy is None and total_enemy is None

        game_state = 0
        if is_start:
            game_state = -1
        elif is_complete:
            game_state = 1

        if show:
            self.show_img(img_i)

        if _img_i is not None:
            self.cache_info_dc.update({_img_i: (game_state, extinct_enemy, total_enemy)})

        return game_state, extinct_enemy, total_enemy

    def _get_is_complete(self, img_i, show=False):
        if isinstance(img_i, int):
            img_i = self.get_img_i(img_i)
        text, score = ocr_utils.get_orc_result(img_i, coordinates=self.coordinates.is_complete_coordinates, show=show)
        res = '已' in text
        return res

    def _get_extinct_enemy_and_total_enemy(self, img_i, show=False):
        if isinstance(img_i, int):
            img_i = self.get_img_i(img_i)
        text, score = ocr_utils.get_orc_result(img_i, coordinates=self.coordinates.extinct_enemy_coordinates, show=show)

        if is_fraction(text):
            extinct_enemy, total_enemy = [int(i) for i in text.split('/')]
        else:
            extinct_enemy, total_enemy = None, None
        return extinct_enemy, total_enemy

    def show_base_df(self, show_path=False):
        columns = self.df.columns if show_path else ['i', 'time', 'file_name']
        print(self.df[columns].head())

    def get_total_enemy(self, debug=False):
        self.total_enemy = self._get_total_enemy_by_binary_search(debug=debug)
        assert self.total_enemy != -1

    def _get_total_enemy_by_binary_search(self, debug=False):

        """
        这里要找到[game_state, extinct_enemy, total_enemy]中的`total_enemy`, 即一个game_state为0的值
        """
        # target = 1
        # arr = self.img_path_ls

        length = self.total_frames

        low, high = 0, length - 1

        target = 0
        while low <= high:
            mid = (low + high) // 2
            game_state, extinct_enemy, total_enemy = self.get_info(mid)
            mid_val = game_state
            if debug:
                print(f'---- low, high: {low, high} --- mid: {mid}, mid_val: {mid_val} --- now: {tt.now()}')
            if mid_val == target:
                if debug:
                    print(f'~~~~~~ mid_val == target --- mid: {mid}, mid_val: {mid_val}')
                return total_enemy  # 找到目标，返回索引
            elif mid_val < target:
                low = mid + 1  # 目标在右半部分
            else:
                high = mid - 1  # 目标在左半部分

        return -1

    def binary_search(self, target_v=0, change_dc=None, range_dc=None):
        """
        这里要找到[game_state, extinct_enemy, total_enemy]中的`total_enemy`, 即一个game_state为0的值
        """
        # target = 1
        # arr = self.img_path_ls

        length = self.total_frames

        self.total_enemy

        if change_dc is None:
            # change_ls = list(range(-1, self.total_enemy + 1))  # start: -1 -> 0, end: 8 -> 9
            change_ls = list(range(self.total_enemy + 1))  # start: -1 -> 0, end: 8 -> 9
            change_dc = dict(zip(change_ls, [-1] * len(change_ls)))
            range_dc = change_dc.copy()

        # range_dc[0] = (0, 100)

        sum(change_dc.values()) == -len(change_ls)  # 全为-1
        list(change_dc.values()).count(-1) == 0  # 不存在-1

        min_i, max_i = 1, length - 1

        change_dc = dict(zip(change_ls, [-1] * len(change_ls)))
        range_dc = change_dc.copy()
        target_dc = change_dc.copy()

        for target_v in change_ls:
        # for target_v in [0, 1, 2]:
        # for target_v in [7]:
            low, high = 1, length - 1
            # target_v = 3
            target_i = -1
            # mid = -1
            # mid_val = -1
            tt.__init__()
            # while low <= high and tt.during(10):
            while low <= high:
                mid = (low + high) // 2

                game_state, extinct_enemy, _ = self.get_info(mid)
                mid_val = self.get_mid_val(game_state, extinct_enemy)

                _game_state, _extinct_enemy, _ = self.get_info(mid - 1)
                _mid_val = self.get_mid_val(_game_state, _extinct_enemy)

                print(f'---- low, high: {low, high} --- mid: {mid}, mid_val: {mid_val} --- now: {tt.now()}')

                if mid_val == target_v:
                    print(f'~~~~~~ mid: {mid}, mid_val: {mid_val}')
                    target_i = mid

                    # high = mid - 1  # 目标在左半部分

                    i = mid_val
                    v = mid
                    if range_dc[i] == -1:
                        range_dc[i] = [min_i, v]
                    else:
                        range_dc[i][1] = min(v, range_dc[i][1])

                    # i = mid
                    # range_dc[i] =
                    # if range_dc[i] == -1:
                    #     range_dc[i] = [low, max_i]
                    # else:
                    #     if low > range_dc[i][0]:
                    #         range_dc[i][0] = low
                    break
                    # return mid  # 找到目标，返回索引
                elif mid_val < target_v:
                    low = mid + 1  # 目标在右半部分
                    # change_dc
                    # range_dc
                    # range_dc = change_dc.copy()

                    for i in range(mid_val + 1, self.total_enemy + 1):
                        if range_dc[i] == -1:
                            range_dc[i] = [low, max_i]
                        else:
                            range_dc[i][0] = max(low, range_dc[i][0])
                            # if low > range_dc[i][0]:
                            #     range_dc[i][0] = low
                else:
                    high = mid - 1  # 目标在左半部分

                    for i in range(0, mid_val):
                        if range_dc[i] == -1:
                            range_dc[i] = [min_i, high]
                        else:
                            range_dc[i][1] = min(high, range_dc[i][1])
                            # if high < range_dc[i][1]:
                            #     range_dc[i][1] = high
                print(f'--- target_v: {target_v} ~~~~~~ range_dc: {range_dc} ====== now: {tt.now()}')

                # print(f'---- low, high: {low, high} --- mid: {mid}, mid_val: {mid_val} --- now: {tt.now()}, target_v: {target_v}')

                # mid_val_0, _, _ = self.get_info(mid)
                # mid_val_1, _, _ = self.get_info(mid + 1)
                # if mid_val_0 == mid_val_1:
                #     if mid_val_0 == 0:

            # assert mid_val == target
            # assert target_i != -1

            if target_i != -1:
                change_dc[target_v] = target_i
                target_dc[target_v] = target_i

            print(f'====== target_v: {target_v}--- range_dc: {range_dc} ====== now: {tt.now()}\n')

        sum([v[1] - v[0] for k, v in range_dc.items()])
        return target_i  # -1则目标不在列表中

        change_dc
        range_dc
        self.get_info(target_i)

        self.cache_info_dc

        index = 659
        self.get_info(index)
        self.get_mid_val(*self.get_info(index)[:2])

        self.get_info(index, show=True)
        self._get_is_complete(index, show=True)

        return target_i  # -1则目标不在列表中

    def get_reward_per_frame(self, batch_size=1, get_rtg=True):
        _reward_state_info = []
        tq_i = tqdm(total=self.total_frames)

        for i in range(0, self.total_frames, batch_size):
            # print(f'i --- {i}')
            game_state, extinct_enemy, _ = self.get_info(i)
            extinct_enemy = extinct_enemy if extinct_enemy else self.total_enemy if game_state == 1 else 0
            _reward_state_info.append([game_state, extinct_enemy, self.total_enemy - extinct_enemy])
            tq_i.update(batch_size)

        self._reward_state_info = _reward_state_info
        df = pd.DataFrame(self._reward_state_info, columns=['game_state', 'extinct_enemy', 'remain_enemy'])
        if get_rtg:
            df = self.get_rtg(df)

        merged_df = pd.concat([self.df, df], axis=1)
        self.reward_state_info_df = self.get_clean_mouse_and_key(merged_df)
        self.reward_state_info_df.to_csv(self.output_path__reward_state_info_df, index=False)

    def get_rtg(self, df, delta__reward_time=None, delta__reward_extinct_enemy=None, delta__reward_change_game_state=None):
        if delta__reward_time is None:
            # delta__reward_time = -1  # 每一帧扣1分
            delta__reward_time = -1 / DefaultArgs.frames  # 每一秒扣1分
            # delta__reward_time = -1 / DefaultArgs.frames / 60  # 每一分钟扣1分

        if delta__reward_extinct_enemy is None:
            delta__reward_extinct_enemy = 1
            # delta__reward_extinct_enemy = 100 / self.total_enemy

        if delta__reward_change_game_state is None:
            delta__reward_change_game_state = 10
            # delta__reward_change_game_state = 100

        df['reward'] = delta__reward_time

        last_extinct_enemy = 0
        last_game_state = -1
        for i, row in df.iterrows():
            # break
            extinct_enemy = row['extinct_enemy']
            game_state = row['game_state']
            reward = row['reward']
            # extinct_enemy = df.loc[i, 'extinct_enemy']
            # game_state = df.loc[i, 'game_state']
            # reward = 0
            if extinct_enemy > last_extinct_enemy:
                reward += delta__reward_extinct_enemy
                print(f'--- extinct_enemy i: {i}')
            if game_state > last_game_state:
                reward += delta__reward_change_game_state
                print(f'--- game_state i: {i}')
            df.loc[i, 'reward'] = reward

            if i < df.shape[0] - 1:
                last_extinct_enemy = extinct_enemy
                last_game_state = game_state

        reward_ls = df['reward'].tolist()

        rtg_ls = []
        rtg_accumulate = 0
        for i in reversed(range(len(reward_ls))):
            # print(i)
            reward_i = reward_ls[i]

            rtg_accumulate += reward_i
            rtg_ls.append(rtg_accumulate)
        rtg_ls.reverse()

        df['rtg'] = rtg_ls

        return df

    def get_clean_mouse_and_key(self, df):
        mouse_pressed_col = []
        key_pressed_col = []

        for i, row in df.iterrows():
            mouse_dc = row['mouse']
            key_dc = row['key']

            mouse_pressed = [k for k, v in mouse_dc.items() if k not in ['time', 'i'] and v != 0]
            mouse_pressed_col.append(mouse_pressed)

            key_pressed = [k for k, v in key_dc.items() if k not in ['time', 'i'] and v != 0]
            key_pressed_col.append(key_pressed)

        df['mouse_pressed'] = pd.Series(mouse_pressed_col)
        df['key_pressed'] = pd.Series(key_pressed_col)
        return df

    def show_video(self, start_i, end_i):
        image_files = self.img_path_ls[start_i: end_i]
        total = end_i - start_i
        tq_i = tqdm(total=total)

        record_time_schedule = DefaultArgs.record_time_schedule[start_i: end_i]

        tt = Time()
        for i in range(len(image_files)):
            image_file = image_files[i]

            next_record_time = record_time_schedule[i] - record_time_schedule[0]

            now = tt.now()
            while now < next_record_time:
                if DefaultArgs.use_sleep_wait:
                    tt.sleep(0.001)
                now = tt.now()

            # msg = f'now: {now}, next_record_time: {next_record_time: .3f}'
            msg = f'delta_t: {now - next_record_time: .3f}'
            print(msg)

            # 读取图片
            # frame = cv2.imread(image_file)
            frame = my_cv2_utils.read_image(image_file)
            # my_cv2_utils.show_img(frame)

            # 显示图片
            cv2.imshow('Frame', frame)

            # 检测键盘按键，如果是 'q' 则退出循环
            if cv2.waitKey(25) & 0xFF == ord('q'):
                break

            tq_i.update(1)
            # tt.sleep(interval)

        # 关闭窗口
        cv2.destroyAllWindows()


def main():
    # img_dir_path = '../tempdir/recorded_states/screen_image/test_7'
    # img_dir_name = 'test_7'
    # operation_name = 'test.json'

    operation_name = 'cao_shen__0'
    recorded_state_dir_path = DefaultArgs.record_path.recorded_state_dir_path

    screen_dir_path = os.path.join(recorded_state_dir_path, 'screen_image', operation_name)
    img_dir_ls = os.listdir(screen_dir_path)
    # img_dir_ls = [i for i in _img_dir_ls if i != '0']
    # img_dir_ls
    for i in range(len(img_dir_ls)):
        # i = 1

        print(f'\n ------ i: {i}')
        tt.sleep(1)

        img_dir_name = img_dir_ls[i]
        record_analysis = RewardAnalysis(operation_name=operation_name, img_dir_name=img_dir_name)
        record_analysis.get_reward_per_frame()
        # break

    exit()

    # from tools.beep import end_beep
    # end_beep()

    img_dir_name = 'cao_shen.json__2024-02-02 08-16-52'
    record_analysis = RewardAnalysis(operation_name=operation_name, img_dir_name=img_dir_name)

    # record_analysis.df.columns
    # df_2 = pd.read_csv(record_analysis.output_path__reward_state_info_df)
    # df_2.head()
    # record_analysis.df.head()
    # record_analysis = RewardAnalysis(operation_name=operation_name, img_dir_name=img_dir_name, i_range=(240, 260))
    # record_analysis.total_frames

    # print(record_analysis.img_dir_path)

    # record_analysis.get_total_enemy(debug=True)
    # record_analysis.total_enemy

    index = record_analysis.key_total_frames

    # duration = 10
    # index = record_analysis.key_total_frames - duration
    # # record_analysis.show_video(index, index + duration)
    # # Coordinates.extinct_enemy_coordinates
    # # record_analysis._get_extinct_enemy_and_total_enemy(260, show=True)
    #
    # game_state, extinct_enemy, total_enemy = record_analysis.get_info(index)
    # print(f"--- index: {index} --- game_state: {game_state} | extinct_enemy, total_enemy: {extinct_enemy, total_enemy}")

    # is_complete = record_analysis._get_is_complete(index)
    # ocr_utils.get_orc_result(record_analysis.get_img_i(index), coordinates=record_analysis.coordinates.is_complete_coordinates)
    # ocr_utils.get_orc_result(record_analysis.get_img_i(index), coordinates=record_analysis.coordinates.extinct_enemy_coordinates)

    # record_analysis.show_img(index)

    record_analysis.get_reward_per_frame()

    pass


if __name__ == '__main__':
    main()






